#include "lib/print.h"
#include "trap/trap.h"
#include "proc/cpu.h"
#include "mem/vmem.h"
#include "memlayout.h"
#include "riscv.h"
#include "trap/intrno.h"
#include "syscall/syscall.h"
#include "mem/pmem.h"
#include "lib/str.h"

// in trampoline.S
extern char trampoline[];  // 内核和用户切换的代码
extern char user_vector[]; // 用户触发trap进入内核
extern char user_return[]; // trap处理完毕返回用户

// in trap.S
extern char kernel_vector[]; // 内核态trap处理流程

// in trap_kernel.c
extern char *interrupt_info[16]; // 中断错误信息
extern char *exception_info[16]; // 异常错误信息

// 判断是否为 cow page
// 是的话返回 1，否则返回 0
int is_cowpage(pgtbl_t pg, uint64 va) {
    if(va >= VA_MAX) 
        return -1;
         
    pte_t *pte = vm_getpte(pg, va, false);
    if (pte == 0 || (*pte & PTE_V) == 0) 
        return 0;
    // 如果页表项有PTE_COW标志，则是COW页面
    return (*pte & PTE_COW) != 0;
}

// handle_cow 处理写时复制的页面错误，分配新页面并复制数据。
void* handle_cow(pgtbl_t pg, uint64 va) {
    if(va % PGSIZE != 0)
        return 0;

    pte_t *pte = vm_getpte(pg, va, false);
    assert(pte != NULL, "handle_cow: pte is NULL");

    uint64 pa = PTE_TO_PA(*pte);
    int flags = PTE_FLAGS(*pte);

    // 如果只存在一个引用，则不需要分配新的物理页
    if(pmem_get_ref(pa) == 1){
        *pte = ((*pte) | PTE_W) & ~PTE_COW;
        return (void*)pa;
    }

    // 存在多个引用，需要分配新的物理页
    char* new_page = (char*)pmem_alloc(false);
    if (new_page == 0) 
        return 0;
        // panic("handle_cow: failed to allocate new page");

    memcpy(new_page, (char *)pa, PGSIZE);

    // 清除PTE_V，否则在mappagges中会判定为remap
    *pte &= ~PTE_V;

    // 为新页面添加映射
    // *pte = PA_TO_PTE((uint64)new_page) | ((flags | PTE_W) & (~PTE_COW));
    vm_mappages(pg, va, (uint64)new_page, PGSIZE, (flags | PTE_W) & ~PTE_COW);
    // if(vm_mappages(pg, va, (uint64)new_page, PGSIZE, (flags | PTE_W) & ~PTE_COW) != 0){
    //     pmem_free((uint64)new_page);
    //     *pte |= PTE_V;
    //     return 0;
    // }

    pmem_free(pa, false);     // 将旧的物理页的 ref_cnt 减1
    return (void *)new_page;
}

// 处理非cow页缺页错误，重新分配新的页
void* handle_page_fault(pgtbl_t pg, uint64 va)
{
    assert (va <= VA_MAX, "handle_page_fault: the va is invalid");
    char* new_page = (char*)pmem_alloc(false);
    memset((void *)new_page, 0, PGSIZE);
    va = PGROUNDDOWN(va);
    vm_mappages(pg, va, (uint64)new_page, PGSIZE, PTE_W | PTE_U | PTE_R);
    return (void *)new_page;
}

// 在user_vector()里面调用
// 用户态trap处理的核心逻辑
void trap_user_handler()
{
    // 这些寄存器保存了陷阱发生时的上下文信息。
    uint64 sepc = r_sepc();       // 记录了发生异常时的pc值
    uint64 sstatus = r_sstatus(); // 与特权模式和中断相关的状态信息
    uint64 scause = r_scause();   // 引发trap的原因
    uint64 stval = r_stval();     // 发生trap时保存的附加信息(不同trap不一样)

    proc_t *p = myproc();
    int trap_id = scause & 0xf;

    // 确认trap来自U-mode
    assert((sstatus & SSTATUS_SPP) == 0, "trap_user_handler: not from u-mode");

    // 将中断和异常发送到 kernel_trap()
    w_stvec((uint64)kernel_vector);

    // 设置 S Exception Program Counter 为保存的用户 PC
    p->tf->epc = sepc;

    if (IS_INTR(scause)){
        switch(trap_id)
        {
            case SMODE_SOFTWARE_INTERRUPT:
                timer_interrupt_handler();
                break;
            case SMODE_EXTERNAL_INTERRUPT:
                external_interrupt_handler();
                break;
            default:
                panic("trap_user_handler:Unhandled interruption,\n\tsepc:%p,scause:%p,sstatus:%p,stval:%p,trap_id:%d\n\tdescription:%s"
                    ,sepc,scause,sstatus,stval,trap_id,interrupt_info[trap_id]);
                break;
        }
    }

    else {
        switch(trap_id)
        {
            case UMODE_SYSCALL_INTERRUPT:
                p->tf->epc += 4;
                intr_on();
                syscall();
                break;

            case STORE_PAGE_FAULTS:
                if(is_cowpage(p->pgtbl, stval) != 0)
                    handle_cow(p->pgtbl, PGROUNDDOWN(stval));
                else 
                    handle_page_fault(p->pgtbl, stval);
                break;

            default:
                panic("trap_user_handler:Unhandled exception,\n\tsepc:%p,scause:%p,sstatus:%p,stval:%p,trap_id:%d\n\tdescription:%s"
                    ,sepc,scause,sstatus,stval,trap_id,exception_info[trap_id]);
                break;
        }
    }

    // 返回用户态
    trap_user_return();
}

// 调用user_return()
// 内核态返回用户态
void trap_user_return()
{
    proc_t* p = myproc();
    intr_off();
    w_stvec(TRAMPOLINE+((uint64)user_vector-(uint64)trampoline));

    // 设置 uservec 需要的陷阱帧值，并恢复用户寄存器
    p->tf->kernel_satp = r_satp();         // 内核页表
    p->tf->kernel_sp = p->kstack + PGSIZE; // 进程的内核栈
    p->tf->kernel_trap = (uint64)trap_user_handler;
    p->tf->kernel_hartid = r_tp(); // hartid 用于 CPU ID

    unsigned long x = r_sstatus();
    x &= ~SSTATUS_SPP; 
    x |= SSTATUS_SPIE;
    w_sstatus(x);

    // 设置 S Exception Program Counter 为保存的用户 PC
    w_sepc(p->tf->epc);

     // 告诉 trampoline.S 切换到的用户页表
    uint64 satp = MAKE_SATP(p->pgtbl);

    // 跳转到 trampoline.S，切换到用户页表，恢复用户寄存器，并通过 sret 切换到用户模式
    uint64 fn = TRAMPOLINE + (user_return - trampoline);
    ((void (*)(uint64, uint64))fn)(TRAPFRAME, satp);
}

//TODO:缺页